import warnings
warnings.filterwarnings('ignore')
from tool import langchain_client
from langchain.prompts import ChatPromptTemplate
from langchain.chains import LLMChain
from langchain.chains import SimpleSequentialChain
if __name__ == '__main__':
    # 提示模板 1 :这个提示将接受产品并返回最佳名称来描述该公司
    first_prompt = ChatPromptTemplate.from_template(
    "描述制造{product}的一个公司的最好的名称是什么，请直接给出名称" )
    chain_one = LLMChain(llm=langchain_client, prompt=first_prompt)
    # 提示模板 2 :接受公司名称，然后输出该公司的长为20个单词的描述
    second_prompt = ChatPromptTemplate.from_template(
    "写一个20字的描述对于下面这个\
    公司:{company_name}的" )
    chain_two = LLMChain(llm=langchain_client, prompt=second_prompt)

    overall_simple_chain = SimpleSequentialChain(chains=[chain_one, chain_two],
                                                 verbose=True)
    product = "大号床单套装"
    response = overall_simple_chain.run(product)
    print(response)